from tqdm import tqdm

from yan import CvFo
import jittor
import pandas as pd
import numpy as np
voc = pd.read_pickle("voc_data.pandas_pickle")
net = CvFo(len(voc), 16, 8)
loss_func = jittor.nn.CrossEntropyLoss()

optimizer = jittor.optim.Adam(net.parameters(), lr=0.0001)
data_set = pd.read_pickle("train_data.pandas_pickle")[:1000]
batch_size = 12
bar = tqdm(len(100 * list(range(0, len(data_set), batch_size))))
loss_list = []
for epoch in range(100):
    np.random.shuffle(data_set)

    for i in range(0, len(data_set), batch_size):
        j = i + batch_size
        one_data = data_set[i:j]
        two_data = jittor.array(one_data).int()

        out= net(two_data[:, :-1])
        loss = loss_func(out.reshape([-1, out.shape[-1]]), two_data[:, 1:].reshape([-1]).long())
        loss_list.append(loss.item())

        bar.set_description("loss___{:.5f}".format(np.mean(loss_list)))
        bar.update(1)
        optimizer.step(loss)

    if (epoch + 1) % 10 == 0:
        jittor.save(net.state_dict(), "model_{}_loss_{:.5f}.pth".format(epoch + 1,np.mean(loss_list)))
jittor.save(net.state_dict(), "model_{}_loss_{:.5f}.pth".format(epoch + 1, np.mean(loss_list)))
